package com.gjxjl.jdbc.gateway.handler;

import com.gjxjl.jdbc.gateway.config.ProxyConfig;
import io.netty.buffer.ByteBuf;
import io.netty.channel.*;
import io.netty.util.ReferenceCountUtil;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.sql.Connection;

public class MySQLProtocolHandler extends ChannelInboundHandlerAdapter {
    private static final Logger logger = LogManager.getLogger(MySQLProtocolHandler.class);

    private final ProxyConfig config;
    private Channel backendChannel;
    private Channel clientChannel;
    private Connection backendConnection;

    public MySQLProtocolHandler(ProxyConfig config) {
        this.config = config;
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        this.clientChannel = ctx.channel();
        logger.info("Client connected: {}", clientChannel.remoteAddress());

        // 在建立连接后添加binlog支持处理器
        ChannelPipeline pipeline = ctx.pipeline();
        pipeline.addAfter(ctx.name(), "binlogSupport", new BinlogSupportHandler(config));

        // 连接到后端MySQL服务器
        connectToBackend(ctx);
    }

    private void connectToBackend(ChannelHandlerContext ctx) {
        try {
            // 这里使用Netty连接到后端MySQL
            BackendConnector.connect(config, ctx.channel(), new BackendConnectionListener() {
                @Override
                public void onBackendConnected(Channel backendChannel) {
                    MySQLProtocolHandler.this.backendChannel = backendChannel;
                    setupChannelForwarding();

                    // 在后端连接上也添加binlog支持
                    backendChannel.pipeline().addFirst("backendBinlogSupport",
                            new BinlogSupportHandler(config));
                }

                @Override
                public void onBackendConnectFailed(Throwable cause) {
                    logger.error("Failed to connect to backend database", cause);
                    ctx.close();
                }
            });
        } catch (Exception e) {
            logger.error("Error connecting to backend", e);
            ctx.close();
        }
    }

    private void setupChannelForwarding() {
        // 设置从前端到后端的转发
        clientChannel.closeFuture().addListener((ChannelFutureListener) future -> {
            if (backendChannel != null && backendChannel.isActive()) {
                backendChannel.close();
            }
        });

        // 设置从后端到前端的转发
        backendChannel.closeFuture().addListener((ChannelFutureListener) future -> {
            if (clientChannel != null && clientChannel.isActive()) {
                clientChannel.close();
            }
        });
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        if (msg instanceof ByteBuf) {
            ByteBuf packet = (ByteBuf) msg;
            try {
                processPacket(ctx, packet);
            } finally {
                // 确保无论发生什么都会释放ByteBuf
                ReferenceCountUtil.release(packet);
            }
        } else {
            // 如果不是ByteBuf，传递给下一个handler
            ctx.fireChannelRead(msg);
        }
    }

    private void processPacket(ChannelHandlerContext ctx, ByteBuf packet) {
        if (config.isEnableLogging()) {
            logPacket("CLIENT -> PROXY", packet);
        }

        // 安全检查
        if (isPacketContainsSQL(packet)) {
            String sql = extractSQLFromPacket(packet);
            if (isSqlBlocked(sql)) {
                logger.warn("Blocked SQL: {}", sql);
                sendErrorResponse(ctx, "SQL blocked by security policy: " + sql);
                return;
            }

            if (config.isEnableLogging()) {
                logger.info("SQL Executed: {}", sql);
            }
        }

        // 转发到后端 - 注意这里需要retain因为我们会释放原始packet
        if (backendChannel != null && backendChannel.isActive()) {
            // retain() 增加引用计数，因为我们要在finally中释放原始packet
            backendChannel.writeAndFlush(packet.retain()).addListener(future -> {
                if (!future.isSuccess()) {
                    logger.error("Failed to forward packet to backend", future.cause());
                }
            });
        } else {
            logger.error("Backend channel is not active");
            ctx.close();
        }
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        logger.info("Client disconnected: {}", ctx.channel().remoteAddress());
        if (backendChannel != null && backendChannel.isActive()) {
            backendChannel.close();
        }
        super.channelInactive(ctx);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        logger.error("Channel exception", cause);
        ctx.close();
    }

    private void logPacket(String direction, ByteBuf packet) {
        try {
            byte[] data = new byte[packet.readableBytes()];
            packet.getBytes(packet.readerIndex(), data);
            logger.debug("{} Packet: {} bytes", direction, data.length);

            // 只记录小包的内容，大包只记录长度
            if (data.length < 100) {
                logger.trace("{} Data: {}", direction, bytesToHex(data));
            }
        } catch (Exception e) {
            logger.warn("Error logging packet", e);
        }
    }

    private boolean isPacketContainsSQL(ByteBuf packet) {
        // 简单判断：COM_QUERY命令是0x03
        if (packet.readableBytes() > 5) {
            int readerIndex = packet.readerIndex();
            try {
                // MySQL包：3字节长度 + 1字节序列号 + 1字节命令
                packet.skipBytes(4); // 跳过长度和序列号
                byte command = packet.readByte();
                return command == 0x03; // COM_QUERY
            } finally {
                packet.readerIndex(readerIndex);
            }
        }
        return false;
    }

    private String extractSQLFromPacket(ByteBuf packet) {
        int readerIndex = packet.readerIndex();
        try {
            if (packet.readableBytes() > 5) {
                packet.skipBytes(5); // 跳过长度、序列号和命令
                byte[] sqlBytes = new byte[packet.readableBytes()];
                packet.readBytes(sqlBytes);
                return new String(sqlBytes).trim();
            }
        } finally {
            packet.readerIndex(readerIndex);
        }
        return "";
    }

    private boolean isSqlBlocked(String sql) {
        if (sql == null || sql.trim().isEmpty()) {
            return false;
        }

        String lowerSql = sql.toLowerCase();
        return config.getBlockedSqlPatterns().stream()
                .anyMatch(pattern -> lowerSql.contains(pattern.toLowerCase()));
    }

    private void sendErrorResponse(ChannelHandlerContext ctx, String error) {
        // 构建MySQL错误响应包
        ByteBuf buffer = ctx.alloc().buffer();

        // 错误包格式: [FF] [错误码] [SQL状态] [错误消息]
        buffer.writeByte(0xFF); // 错误标志
        buffer.writeShortLE(1064); // 错误码
        buffer.writeBytes("#HY000".getBytes()); // SQL状态
        buffer.writeBytes(error.getBytes()); // 错误消息

        ctx.writeAndFlush(buffer);
    }

    private String bytesToHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(String.format("%02x ", b));
        }
        return sb.toString();
    }
}
